{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analysis of LAMMPS MD trajectory:  Preparation of VASP validation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lmpdump as lmpdump    # Use lmpdump.py parser\n",
    "import os\n",
    "import numpy as np\n",
    "import math\n",
    "import itertools\n",
    "\n",
    "pwd = os.getcwd()\n",
    "\n",
    "print('This script extracts frames from LAMMPS NVT trajectory and generates VASP POSCAR files for DFT validation of the atomic forces.')\n",
    "print('This script is interactive and requires user input. Please read carefully before proceeding:\\n')\n",
    "\n",
    "print('Note 1: Intended only for VASP DFT evaluation of the atomic forces (SCF calculation).')\n",
    "print('Note 2: Requires a LAMMPS DATA file containing the equilibrated box dimensions.')\n",
    "print('Note 3: Requires a trajectory containing all force components.')\n",
    "print('Note 4: Requires lmpdump.py to load the LAMMPS trajectories.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load LAMMPS trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LAMMPS custom-style dump file - Raw trajectory\n",
    "name = input('Enter the name of the LAMMPS trajectory file (ID, type, x, y, z, fx, fy, fz) : \\nWarning: Must contain all force components! Please quit now if not so.\\n')\n",
    "print('Loading trajectory...')\n",
    "file = pwd + '/' + name\n",
    "xyzf = lmpdump.lmpdump(file, loadmode='all')\n",
    "print('\\n')\n",
    "\n",
    "print('Loading time steps...')\n",
    "step = []\n",
    "for s in xyzf.finaldict.keys():    # keys = time steps\n",
    "    step.append(s)\n",
    "print('Done.\\n')\n",
    "\n",
    "size = np.array(step).size    # Number of time steps\n",
    "N_dump = xyzf.finaldict[0][1].shape[0]    # Number of atoms dumped"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract unit cell & system information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load LAMMPS data file after NPT equilibration\n",
    "name = input('Enter the name of the LAMMPS DATA file containing the equilibrated box dimensions : ')\n",
    "file_eq = pwd + '/' + name\n",
    "log = open(file_eq,'r')\n",
    "log_lines = log.readlines()\n",
    "log_lines = [line.split() for line in log_lines]\n",
    "\n",
    "for line_index, line in enumerate(log_lines):\n",
    "    if line:\n",
    "        for l in line:\n",
    "            \n",
    "            if l == 'atoms':\n",
    "                N_atom = int(line[0])\n",
    "                \n",
    "            elif l == 'types':\n",
    "                N_element = int(line[0])\n",
    "\n",
    "            elif l == 'xlo':\n",
    "                xlo = float(line[0])\n",
    "                xhi = float(line[1])\n",
    "            \n",
    "            elif l == 'ylo':\n",
    "                ylo = float(line[0])\n",
    "                yhi = float(line[1])\n",
    "            \n",
    "            elif l == 'zlo':\n",
    "                zlo = float(line[0])\n",
    "                zhi = float(line[1])\n",
    "            \n",
    "            elif l == 'xy':\n",
    "                xy = float(line[0])\n",
    "                xz = float(line[1])\n",
    "                yz = float(line[2])\n",
    "            \n",
    "            elif l == 'Atoms':\n",
    "                head = line_index + 1\n",
    "                start = line_index + 2\n",
    "            \n",
    "            elif l == 'Velocities':\n",
    "                end = line_index - 1\n",
    "\n",
    "Lx = xhi - xlo    # Box dimensions\n",
    "Ly = yhi - ylo\n",
    "Lz = zhi - zlo\n",
    "tan = Ly/xy\n",
    "\n",
    "lat = np.array([[xhi-xlo, 0.0, 0.0], [xy, yhi-ylo, 0.0], [xz, yz, zhi-zlo]])    # LAMMPS NVT lattice matrix\n",
    "\n",
    "for line in log_lines[start:end]:\n",
    "    line[:2] = np.array(line[:2]).astype(int)    # Atom ID, atom type\n",
    "    line[2:5] = np.array(line[2:5]).astype(float)    # x, y, z\n",
    "    line[5:8] = np.array(line[5:8]).astype(int)    # ix, iy, iz (inverse sign)\n",
    "\n",
    "fix = (N_dump != N_atom)    # Is bottommost layer muted?\n",
    "if fix:\n",
    "    \n",
    "    print('Fixed coordinates were muted. They will be obtained from the LAMMPS DATA file.')\n",
    "    print('The fixed atoms are assumed to be numbered consecutively (e.g. layer-by-layer). Please quit now if not so.')\n",
    "    fix_range = input('Enter the IDs of the first and last atom to be fixed, separated by space : ')\n",
    "    fix_range = fix_range.split()\n",
    "    fix_range = np.array(fix_range).astype(int)\n",
    "    \n",
    "sys = input('Enter the system name : ')\n",
    "\n",
    "scale = 1.0    # Lattice scaling factor\n",
    "\n",
    "ls_element = input('Enter the list of the elements, separated by space, in the same order as LAMMPS atom types : ')\n",
    "\n",
    "N_type = []    # List of number of atoms of each atom type\n",
    "for atom_type in range(1,N_element+1):\n",
    "    N = 0\n",
    "    for line in log_lines[start:end]:\n",
    "        Type = line[1]\n",
    "        if Type == atom_type:\n",
    "            N += 1\n",
    "    N_type.append(N)\n",
    "N_type = ' '.join(str(n) for n in N_type)\n",
    "\n",
    "header = '%s\\n%f\\n%f %f %f\\n%f %f %f\\n%f %f %f\\n%s\\n%s\\n%s\\n'%(sys, scale, lat[0][0], lat[0][1], lat[0][2], lat[1][0], lat[1][1], lat[1][2], lat[2][0], lat[2][1], lat[2][2], ls_element, N_type, 'Direct')\n",
    "\n",
    "log.close()\n",
    "\n",
    "print('Unit cell & system information extracted.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract snapshots - VASP POSCAR preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_frm = int(input('Enter the number of frames to be extracted in regular step interval, starting with the first frame : '))\n",
    "d = math.floor(size/(N_frm-1))    # Grab interval\n",
    "\n",
    "for n in range(0,N_frm):\n",
    "    \n",
    "    s = step[n*d]\n",
    "    \n",
    "    frm_dir = pwd + '/frm' + str(n+1)\n",
    "    os.mkdir(frm_dir)\n",
    "    \n",
    "    file_poscar = frm_dir + '/POSCAR'\n",
    "    poscar = open(file_poscar,'w')\n",
    "    poscar.write(header)\n",
    "    \n",
    "    for atom_type in range(1,N_element+1):    # List atoms element-by-element\n",
    "    \n",
    "        if fix:    # Obtain fixed atoms from LAMMPS DATA file, if any\n",
    "            for line in log_lines[start:end]:\n",
    "                \n",
    "                ID = line[0]\n",
    "                Type = line[1]\n",
    "                xu = line[2]\n",
    "                yu = line[3]\n",
    "                zu = line[4]\n",
    "                ix = line[5]\n",
    "                iy = line[6]\n",
    "                iz = line[7]\n",
    "                \n",
    "                if (Type == atom_type) and (fix_range[0] <= ID <= fix_range[1]):\n",
    "                    \n",
    "                    x = xu + ix*Lx + iy*xy    # Wrapped coordinates\n",
    "                    y = yu + iy*Ly\n",
    "                    z = zu + iz*Lz\n",
    "                    r = np.array([x, y, z])    # Cartesian coordinates\n",
    "                    Rep = np.dot(r, np.linalg.inv(lat))    # Direct coordinates\n",
    "                    \n",
    "                    coord = '%f\\t%f\\t%f\\n'%(Rep[0], Rep[1], Rep[2])\n",
    "                    poscar.write(coord)\n",
    "        \n",
    "        for atom in range(0,N_dump):    # Obtain mobile atoms from LAMMPS trajectory\n",
    "            \n",
    "            ID = xyzf.finaldict[s][1]['id'][atom]\n",
    "            Type = xyzf.finaldict[s][1]['type'][atom]\n",
    "            \n",
    "            if Type == atom_type:\n",
    "                \n",
    "                x = xyzf.finaldict[s][1]['x'][atom]\n",
    "                y = xyzf.finaldict[s][1]['y'][atom]\n",
    "                z = xyzf.finaldict[s][1]['z'][atom]\n",
    "                r = np.array([x, y, z])\n",
    "                Rep = np.dot(r, np.linalg.inv(lat))\n",
    "                \n",
    "                coord = '%f\\t%f\\t%f\\n'%(Rep[0], Rep[1], Rep[2])\n",
    "                poscar.write(coord)\n",
    "    \n",
    "    poscar.close()\n",
    "\n",
    "print('POSCAR generated > ./validation/frm**/POSCAR')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract force components after calculations completed: Force field vs. DFT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_frm = int(input('Enter the number of frames extracted : '))\n",
    "d = math.floor(size/(N_frm-1))    # Grab interval\n",
    "\n",
    "for n in range(0,N_frm):\n",
    "    \n",
    "    s = step[n*d]    # Time step\n",
    "    \n",
    "    frm_dir = pwd + '/frm' + str(n+1)\n",
    "    file_out = frm_dir + '/xyzf.out'    # Position & force data from DFT\n",
    "    out = open(file_out,'r')\n",
    "    out_lines = out.readlines()\n",
    "    out_lines = [line.split() for line in out_lines]\n",
    "    \n",
    "    for line_index, line in enumerate(out_lines):\n",
    "        if line_index > 1:\n",
    "            line[:] = np.array(line[:]).astype(float)    # x, y, z, fx, fy, fz\n",
    "        \n",
    "    index = 2\n",
    "    for atom_type in range(1,N_element+1):    # Atoms listed element-by-element\n",
    "        \n",
    "        # List of force components to be plotted for each element\n",
    "        file_force = pwd + '/f_validate-' + str(atom_type) + '.txt'\n",
    "        force = open(file_force,'a')\n",
    "        \n",
    "        if n == 0:\n",
    "            header = 'Force field\\tDFT\\n'\n",
    "            force.write(header)\n",
    "\n",
    "        if fix:    # Skip fixed atoms listed first (zero forces in LAMMPS)\n",
    "            for line in log_lines[start:end]:\n",
    "                \n",
    "                ID = line[0]\n",
    "                Type = line[1]\n",
    "                \n",
    "                if (Type == atom_type) and (fix_range[0] <= ID <= fix_range[1]):\n",
    "                    index += 1\n",
    "                    \n",
    "        for atom in range(0,N_dump):\n",
    "            \n",
    "            Type = xyzf.finaldict[s][1]['type'][atom]\n",
    "            \n",
    "            if Type == atom_type:\n",
    "                \n",
    "                fx_ff = xyzf.finaldict[s][1]['fx'][atom]    # Force components from force field\n",
    "                fy_ff = xyzf.finaldict[s][1]['fy'][atom]\n",
    "                fz_ff = xyzf.finaldict[s][1]['fz'][atom]\n",
    "                \n",
    "                fx_dft = out_lines[index][3]    # Force components from DFT\n",
    "                fy_dft = out_lines[index][4]\n",
    "                fz_dft = out_lines[index][5]\n",
    "                \n",
    "                index += 1\n",
    "                \n",
    "                line_x = '%f\\t%f\\n'%(fx_ff, fx_dft)\n",
    "                line_y = '%f\\t%f\\n'%(fy_ff, fy_dft)\n",
    "                line_z = '%f\\t%f\\n'%(fz_ff, fz_dft)\n",
    "                \n",
    "                force.write(line_x)\n",
    "                force.write(line_y)\n",
    "                force.write(line_z)\n",
    "        \n",
    "        force.close()\n",
    "        \n",
    "    out.close()\n",
    "\n",
    "print('Force field vs. DFT force components extracted for each element > ./f_validate-*.txt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate errors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwd = os.getcwd()\n",
    "\n",
    "file_ls = os.popen('ls -1 f_validate-*.txt').readlines()    # File list of force components for each element\n",
    "file_ls = [line.split() for line in file_ls]\n",
    "file_ls = list(itertools.chain.from_iterable(file_ls))\n",
    "\n",
    "file_error = pwd + '/error.txt'\n",
    "error = open(file_error,'w')\n",
    "\n",
    "for file_index, file in enumerate(file_ls):\n",
    "    \n",
    "    file_dir = pwd + '/' + file\n",
    "    force = open(file_dir,'r')\n",
    "    force_lines = force.readlines()\n",
    "    force_lines = [line.split() for line in force_lines]\n",
    "    \n",
    "    RMSE = 0.0    # Root mean squre error\n",
    "    MAE = 0.0    # Mean absolute error\n",
    "    for line in force_lines[1:]:\n",
    "        line[:] = np.array(line[:]).astype(float)\n",
    "        RMSE += (line[0]-line[1])**2\n",
    "        MAE += abs(line[0]-line[1])\n",
    "    RMSE = np.sqrt(RMSE / len(force_lines[1:]))\n",
    "    MAE = MAE / len(force_lines[1:])\n",
    "\n",
    "    Type = file_index + 1\n",
    "    header = 'Element #' + str(Type) + '\\n'\n",
    "    out_RMSE = '%s%f%s\\n'%('RMSE = ', RMSE, ' eV/Ang')\n",
    "    out_MAE = '%s%f%s\\n'%('MAE = ', MAE, ' eV/Ang')\n",
    "    \n",
    "    error.write(header)\n",
    "    error.write(out_RMSE)\n",
    "    error.write(out_MAE)\n",
    "    error.write('\\n')\n",
    "\n",
    "error.close()\n",
    "\n",
    "print('Error in force components calculated for each element > ./error.txt')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
